"""
背景移除核心功能模块
"""

import os
import sys
from pathlib import Path
from typing import Union, List, Optional

try:
    import numpy as np
except ImportError as e:
    print("❌ 缺少必要依赖，请安装: pip install numpy")
    sys.exit(1)

try:
    from PIL import Image
except ImportError as e:
    print("❌ 缺少必要依赖，请安装: pip install Pillow")
    sys.exit(1)

try:
    import cv2
except ImportError as e:
    print("❌ 缺少必要依赖，请安装: pip install opencv-python")
    sys.exit(1)

try:
    from rembg import remove, new_session
except ImportError as e:
    print("❌ 缺少必要依赖，请安装: pip install rembg")
    sys.exit(1)


class BackgroundRemover:
    """背景移除器"""
    
    def __init__(self, model_name: str = "u2net"):
        """
        初始化背景移除器
        
        Args:
            model_name: 使用的AI模型名称 (u2net, u2net_human_seg, silueta等)
        """
        self.model_name = model_name
        self.session = None
        self._init_session()
    
    def _init_session(self):
        """初始化模型会话"""
        try:
            print(f"[初始化] 初始化AI模型: {self.model_name}")
            
            # 设置模型路径环境变量
            self._setup_model_path()
            
            # 初始化模型会话
            self.session = new_session(self.model_name)
            print("[成功] 模型加载成功")
            
        except Exception as e:
            print(f"[错误] 模型加载失败: {e}")
            print("[建议] 请检查网络连接或运行 python fix_network_error.py")
            raise
    
    def _setup_model_path(self):
        """设置模型路径环境变量"""
        # 检查是否在打包的exe中
        if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
            # 如果是打包后的exe文件
            bundled_models_dir = Path(sys._MEIPASS) / 'models'
            if bundled_models_dir.exists():
                os.environ['U2NET_HOME'] = str(bundled_models_dir)
                print(f"[模型路径] 使用打包的模型: {bundled_models_dir}")
                return
        
        # 检查当前目录下的models文件夹（开发模式）
        local_models_dir = Path.cwd() / 'models'
        if local_models_dir.exists() and any(local_models_dir.glob('*.onnx')):
            os.environ['U2NET_HOME'] = str(local_models_dir)
            print(f"[模型路径] 使用本地模型: {local_models_dir}")
            return
    
    def remove_background(self, input_path: Union[str, Path], 
                         output_path: Optional[Union[str, Path]] = None,
                         output_format: str = "PNG") -> Path:
        """
        移除单张图片的背景
        
        Args:
            input_path: 输入图片路径
            output_path: 输出图片路径（可选）
            output_format: 输出格式 (PNG, JPEG等)
            
        Returns:
            输出文件路径
        """
        input_path = Path(input_path)
        
        if not input_path.exists():
            raise FileNotFoundError(f"输入文件不存在: {input_path}")
        
        # 生成输出路径
        if output_path is None:
            output_path = input_path.parent / f"{input_path.stem}_no_bg.png"
        else:
            output_path = Path(output_path)
        
        try:
            print(f"[处理] 处理图片: {input_path.name}")
            
            # 读取输入图片
            with open(input_path, 'rb') as f:
                input_data = f.read()
            
            # 移除背景
            output_data = remove(input_data, session=self.session)
            
            # 保存结果
            try:
                output_path.parent.mkdir(parents=True, exist_ok=True)
            except PermissionError:
                # 如果没有权限，优先使用桌面目录
                desktop_dir = Path.home() / "Desktop" / "抠图输出"
                desktop_dir.mkdir(parents=True, exist_ok=True)
                output_path = desktop_dir / output_path.name
                print(f"[权限] 输出到桌面目录: {output_path}")
            
            with open(output_path, 'wb') as f:
                f.write(output_data)
            
            print(f"[成功] 完成: {output_path.name}")
            return output_path
            
        except Exception as e:
            print(f"[错误] 处理失败 {input_path.name}: {e}")
            raise
    
    def batch_remove_background(self, input_dir: Union[str, Path],
                              output_dir: Optional[Union[str, Path]] = None,
                              file_extensions: Optional[List[str]] = None) -> List[Path]:
        """
        批量移除背景
        
        Args:
            input_dir: 输入目录
            output_dir: 输出目录（可选）
            file_extensions: 支持的文件扩展名
            
        Returns:
            处理成功的文件路径列表
        """
        input_dir = Path(input_dir)
        
        if file_extensions is None:
            file_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
        
        if output_dir is None:
            output_dir = input_dir / "no_background"
        else:
            output_dir = Path(output_dir)
        
        # 尝试创建输出目录，如果失败则使用备用方案
        try:
            output_dir.mkdir(parents=True, exist_ok=True)
            print(f"[目录] 输出目录: {output_dir}")
        except PermissionError:
            # 如果没有权限，优先使用桌面目录
            desktop_dir = Path.home() / "Desktop" / "抠图输出"
            try:
                desktop_dir.mkdir(parents=True, exist_ok=True)
                output_dir = desktop_dir
                print(f"[权限] 无法在原目录创建输出文件夹，已在桌面创建: {output_dir}")
            except Exception as e2:
                # 如果桌面也失败，尝试文档目录
                documents_dir = Path.home() / "Documents" / "抠图输出"
                try:
                    documents_dir.mkdir(parents=True, exist_ok=True)
                    output_dir = documents_dir
                    print(f"[权限] 桌面创建失败，使用文档目录: {output_dir}")
                except Exception as e3:
                    print(f"[错误] 无法创建任何输出目录")
                    print(f"  - 原目录错误: 没有写权限")
                    print(f"  - 桌面目录错误: {e2}")
                    print(f"  - 文档目录错误: {e3}")
                    print(f"[建议] 请将程序移动到有写权限的目录（如桌面或文档文件夹）再运行")
                    raise PermissionError("无法创建输出目录，请检查文件权限")
        except Exception as e:
            print(f"[错误] 创建输出目录失败: {e}")
            raise
        
        # 查找所有图片文件
        image_files = []
        for ext in file_extensions:
            image_files.extend(input_dir.glob(f"*{ext}"))
            image_files.extend(input_dir.glob(f"*{ext.upper()}"))
        
        if not image_files:
            print(f"[警告] 在目录 {input_dir} 中未找到支持的图片文件")
            return []
        
        print(f"[文件] 找到 {len(image_files)} 个图片文件")
        
        success_files = []
        failed_files = []
        
        for i, img_file in enumerate(image_files, 1):
            try:
                output_path = output_dir / f"{img_file.stem}_no_bg.png"
                print(f"[{i}/{len(image_files)}] ", end="")
                
                result_path = self.remove_background(img_file, output_path)
                success_files.append(result_path)
                
            except Exception as e:
                print(f"[错误] 跳过文件 {img_file.name}: {e}")
                failed_files.append(img_file)
        
        print(f"\n[统计] 批处理完成:")
        print(f"   [成功] 成功: {len(success_files)} 个文件")
        print(f"   [失败] 失败: {len(failed_files)} 个文件")
        
        if success_files:
            print(f"   [目录] 输出目录: {output_dir}")
        
        return success_files
    
    def add_new_background(self, no_bg_image_path: Union[str, Path], 
                          background_path: Union[str, Path],
                          output_path: Optional[Union[str, Path]] = None) -> Path:
        """
        为已抠图的图片添加新背景
        
        Args:
            no_bg_image_path: 无背景图片路径
            background_path: 新背景图片路径
            output_path: 输出路径
            
        Returns:
            输出文件路径
        """
        no_bg_path = Path(no_bg_image_path)
        bg_path = Path(background_path)
        
        if output_path is None:
            output_path = no_bg_path.parent / f"{no_bg_path.stem}_new_bg.png"
        else:
            output_path = Path(output_path)
        
        try:
            # 读取图片
            foreground = Image.open(no_bg_path).convert("RGBA")
            background = Image.open(bg_path).convert("RGB")
            
            # 调整背景尺寸匹配前景
            background = background.resize(foreground.size, Image.Resampling.LANCZOS)
            background = background.convert("RGBA")
            
            # 合成图片
            result = Image.alpha_composite(background, foreground)
            result = result.convert("RGB")
            
            # 保存结果
            try:
                output_path.parent.mkdir(parents=True, exist_ok=True)
            except PermissionError:
                # 如果没有权限，优先使用桌面目录
                desktop_dir = Path.home() / "Desktop" / "抠图输出"
                desktop_dir.mkdir(parents=True, exist_ok=True)
                output_path = desktop_dir / output_path.name
                print(f"[权限] 输出到桌面目录: {output_path}")
            
            result.save(output_path, "PNG")
            
            print(f"[成功] 背景替换完成: {output_path.name}")
            return output_path
            
        except Exception as e:
            print(f"[错误] 背景替换失败: {e}")
            raise